include("BoltzmannMachine.jl")
include("Utils.jl")


"""
Implements the theorem "Perturbation Solutions of BM".
"""
function perturbsolve(x̂::AbstractVector{T}, Ĉ::AbstractMatrix{T}) where T<:Real
    # TODO: Add validation of the perturbation condition.

    W = @. zero(Ĉ)
    @inbounds for i = 1:size(W, 1)
        for j = 1:size(W, 2)
            if i != j
                W[i, j] = Ĉ[i, j] / x̂[i] / (1 - x̂[i]) / x̂[j] / (1 - x̂[j])
            end
        end
    end

    c = @. invσ(x̂)
    b = c .- W * x̂

    create_network(W, b)
end


"""
Implements the theorem "Perturbation Solutions of BM".
"""
function perturbsolve(X::AbstractMatrix{T}) where T<:Real
    x̂ = expect(X)
    Ĉ = covariance(X)
    perturbsolve(x̂, Ĉ)
end


"""
Implements the "Zoom-in Trick".
"""
function zoomin(m::I, x::T) where {I<:Integer, T<:Real}
    [sample_bernoulli(x) for i = 1:m]
end


"""
Implements the "Zoom-in Trick".
"""
function zoomin(m::AbstractVector{I}, x::AbstractVector{T}) where {I<:Integer, T<:Real}
    vcat([zoomin(mᵢ, xᵢ) for (mᵢ, xᵢ) in zip(m, x)]...)
end


"""
Implements the "Zoom-in Trick".
"""
function zoomin(m::I, x::AbstractVector{T}) where {I<:Integer, T<:Real}
    m_arr = [m for i = 1:size(x, 1)]
    zoomin(m_arr, x)
end


"""
Implements the "Zoom-in Trick".
"""
function zoomin(m::I, x::AbstractMatrix{T}) where {I<:Integer, T<:Real}
    hcat([zoomin(m, x[:, i]) for i = 1:size(x, 2)]...)
end


"""
Inverse of `zoomin`.
"""
function zoomout(m::AbstractVector{I}, x::AbstractVector{T}) where {I<:Integer, T<:Real}
    N = size(x, 1)
    @assert sum(m) == N

    y = zeros(eltype(x), size(m, 1))
    offset = 0
    for i = 1:size(m, 1)
        piece = x[(offset + 1):(offset + m[i])]
        y[i] = mean(piece)
        offset += m[i]
    end
    y
end



"""
Inverse of `zoomin`.
"""
function zoomout(m::I, x::AbstractVector{T}) where {I<:Integer, T<:Real}
    m_arr = [m for i = 1:Int(size(x, 1)/m)]
    zoomout(m_arr, x)
end


"""
Inverse of `zoomin`.
"""
function zoomout(m::I, x::AbstractMatrix{T}) where {I<:Integer, T<:Real}
    hcat([zoomout(m, x[:, i]) for i = 1:size(x, 2)]...)
end
